import gym
import env
import random
import hydra
from env.tsp_parallel import TSPParallelEnv

@ hydra.main(config_path='./', config_name='local_gpu_config.yaml')
def main(config_args):
    # env = TSP_reverse()
    # envs = TSPParallelEnv(config_args.rl_script_args)
    # envs.reset()
    # action_space = config_args.rl_script_args["action_space"]
    # done = False
    # while not done:
    #     action = [random.randint(0, 44) for _ in range(config_args.rl_script_args["number_envs"])]
    #     obss, rewards, dones, infos = envs.step(action)
    #     done = False
    #     for i,d in enumerate(dones):
    #         done = done or d
    #         if i==0:
    #             print(f"env:{i}, obs:{infos[i]['descriptions']}, reward:{rewards[i]}, done:{dones[i]}, action:{action_space[action[i]]}\n")   
    envs = gym.make("TSP_reverse-v0")
    envs.reset()
    action_space = config_args.rl_script_args["action_space"]
    dones = False
    while not dones:
        action = random.randint(0, 44)
        obss, rewards, dones, _, infos = envs.step(action)
        print(f"obs:{infos['descriptions']}, reward:{rewards}, done:{dones}, action:{action_space[action]}\n")  
if __name__ == '__main__':
    main()